import json
import sys
import timeit
import types
from pathlib import Path
from typing import Any

import equinox as eqx
from equinox._module._prebuilt import BoundMethod


# ---------------------------------------------------------------------------
# Module definitions
# ---------------------------------------------------------------------------


class _Simple(eqx.Module):
    a: int
    b: float
    c: str


class _WithConverter(eqx.Module):
    x: int = eqx.field(converter=int)


class _WithCheckInit(eqx.Module):
    a: int

    def __check_init__(self):
        pass


class _DeepA(eqx.Module):
    a: int

    def __check_init__(self):
        pass


class _DeepB(_DeepA):
    def __check_init__(self):
        pass


class _DeepC(_DeepB):
    def __check_init__(self):
        pass


class _ManyFields(eqx.Module):
    f0: Any
    f1: Any
    f2: Any
    f3: Any
    f4: Any
    f5: Any
    f6: Any
    f7: Any
    f8: Any
    f9: Any


class _StaticField(eqx.Module):
    name: str = eqx.field(static=True)
    value: int


class _Wrapped(eqx.Module):
    fn: Any

    def __call__(self, *args, **kwargs):
        return self.fn(*args, **kwargs)

    @property
    def __wrapped__(self):
        return self.fn


def _make_wrapped():
    return eqx.module_update_wrapper(_Wrapped(lambda x: x))


class _Host(eqx.Module):
    x: float
    y: float

    def method(self):
        return self.x + self.y

    def method2(self):
        return self.x * self.y

    def method3(self):
        return self.x - self.y


_host = _Host(1.0, 2.0)
_func: types.FunctionType = _Host.method  # type: ignore[assignment]


# ---------------------------------------------------------------------------
# Benchmark runner
# ---------------------------------------------------------------------------

SCENARIOS: list[tuple[str, str, dict[str, Any]]] = [
    ("simple", "_Simple(1, 2.0, 'x')", {}),
    ("converter", "_WithConverter(42)", {}),
    ("check_init", "_WithCheckInit(7)", {}),
    ("deep_inherit", "_DeepC(3)", {}),
    (
        "many_fields",
        "_ManyFields(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)",
        {},
    ),
    ("static_field", "_StaticField('hello', 99)", {}),
    ("module_update_wrapper", "_make_wrapped()", {}),
    (
        "BoundMethod_init",
        "BoundMethod(_func, _host)",
        {"BoundMethod": BoundMethod, "_func": _func, "_host": _host},
    ),
    (
        "attr_access_single",
        "_host.method",
        {"_host": _host},
    ),
    (
        "attr_access_three",
        "_host.method; _host.method2; _host.method3",
        {"_host": _host},
    ),
]

NUMBER = 200_000
REPEAT = 5


def run_benchmarks() -> dict[str, float]:
    g = {
        "_Simple": _Simple,
        "_WithConverter": _WithConverter,
        "_WithCheckInit": _WithCheckInit,
        "_DeepC": _DeepC,
        "_ManyFields": _ManyFields,
        "_StaticField": _StaticField,
        "_make_wrapped": _make_wrapped,
        "_host": _host,
        "_func": _func,
        "BoundMethod": BoundMethod,
    }
    results: dict[str, float] = {}
    for name, stmt, extra in SCENARIOS:
        times = timeit.repeat(
            stmt, globals={**g, **extra}, number=NUMBER, repeat=REPEAT
        )
        # µs per call
        results[name] = min(times) / NUMBER * 1e6
    return results


if __name__ == "__main__":
    out_path = Path(sys.argv[1]) if len(sys.argv) > 1 else None

    print(f"equinox version: {eqx.__version__}")
    print(f"number={NUMBER:,}  repeat={REPEAT}  (reporting min-of-repeats / call)\n")

    results = run_benchmarks()

    for name, us in results.items():
        print(f"  {name:<20s} {us:.4f} µs/call")

    if out_path is not None:
        out_path.write_text(json.dumps(results, indent=2))
        print(f"\nResults saved to {out_path}")